import os
import json
from django.http import JsonResponse
from django.views.decorators.csrf import csrf_exempt
from paddleocr import PaddleOCR
from PIL import Image
import numpy as np
from rest_framework.decorators import api_view, permission_classes
from rest_framework.permissions import AllowAny
from pdf2image import convert_from_path  # 用于将 PDF 转换为图片
import docx  # 用于处理 Word 文件
import fitz  # 或 from PyMuPDF import fitz
import datetime
import cv2
from paddleocr import PPStructure, save_structure_res
from paddleocr.ppstructure.recovery.recovery_to_doc import sorted_layout_boxes, convert_info_docx
from copy import deepcopy
# 初始化 OCR 模型
ocr = PaddleOCR(use_angle_cls=True, lang='ch')  # 语言设置为中文
def extract_text(line):
    if isinstance(line[1][0], str):
        return line[1][0]
    return ""
def pdf2png(pdfPath, baseImagePath):
    imagePath = os.path.join(baseImagePath, os.path.basename(pdfPath).split('.')[0])
    startTime_pdf2img = datetime.datetime.now()
    print("正在处理PDF文件，保存路径：" + imagePath)
    
    if not os.path.exists(imagePath):
        os.makedirs(imagePath)
    
    try:
        pdfDoc = fitz.open(pdfPath)
        totalPage = pdfDoc.pageCount
        
        for pg in range(totalPage):
            page = pdfDoc[pg]
            rotate = int(0)
            # 增加分辨率以提高清晰度
            zoom_x = 4
            zoom_y = 4
            mat = fitz.Matrix(zoom_x, zoom_y).prerotate(rotate)
            pix = page.get_pixmap(matrix=mat, alpha=False)
            
            # 优化图像质量
            output_path = os.path.join(imagePath, f'page_{pg+1}.png')
            pix.save(output_path)
            print(f'已完成第{pg+1}页转换，共{totalPage}页')
        
        endTime_pdf2img = datetime.datetime.now()
        print(f'PDF转图片完成，耗时：{(endTime_pdf2img - startTime_pdf2img).seconds}秒')
        return imagePath
    except Exception as e:
        print(f'PDF转换失败：{str(e)}')
        raise

@api_view(['POST'])
@permission_classes([AllowAny])
def ocr_detect(request):
    if request.method == 'POST':
        file = request.FILES.get('file')
        if not file:
            return JsonResponse({'error': '未上传文件'}, status=400)

        # 创建media目录（如果不存在）
        media_dir = 'media'
        if not os.path.exists(media_dir):
            os.makedirs(media_dir)

        file_path = os.path.join(media_dir, file.name)
        with open(file_path, 'wb') as f:
            f.write(file.read())

        texts = []
        file_extension = file.name.split('.')[-1].lower()

        try:
            if file_extension in ['jpg', 'jpeg', 'png', 'bmp', 'gif']:
                result = ocr.ocr(file_path, cls=True)
                if result:
                    for line in result:
                        if isinstance(line, list) and len(line) > 0:
                            text = line[1][0] if isinstance(line[1], tuple) else ''
                            if text:
                                texts.append(text)

            elif file_extension == 'pdf':
                pdf_images_dir = os.path.join(media_dir, 'pdf_images')
                if not os.path.exists(pdf_images_dir):
                    os.makedirs(pdf_images_dir)

                image_folder = pdf2png(file_path, pdf_images_dir)
                if image_folder:
                    for img_name in sorted(os.listdir(image_folder)):
                        if img_name.endswith('.png'):
                            img_path = os.path.join(image_folder, img_name)
                            result = ocr.ocr(img_path, cls=True)
                            if result:
                                for line in result:
                                    if isinstance(line, list) and len(line) > 0:
                                        text = line[1][0] if isinstance(line[1], tuple) else ''
                                        if text:
                                            texts.append(text)
                                texts.append('\n')

            elif file_extension in ['doc', 'docx']:
                doc = docx.Document(file_path)
                texts = [para.text for para in doc.paragraphs if para.text.strip()]

            else:
                return JsonResponse({'error': '不支持的文件格式'}, status=400)

            if not texts:
                return JsonResponse({'error': '未检测到文本内容'}, status=404)

            # 清理临时文件
            if os.path.exists(file_path):
                os.remove(file_path)

            return JsonResponse({'text': '\n'.join(texts)})

        except Exception as e:
            if os.path.exists(file_path):
                os.remove(file_path)
            return JsonResponse({'error': f'处理文件时出错: {str(e)}'}, status=500)

    return JsonResponse({'error': '无效的请求方法'}, status=405)
def img2docx(img_path):
    text = []
    imgs = os.listdir(img_path)
    for img_name in imgs:
        print(os.path.join(img_path, img_name))
        img = cv2.imread(os.path.join(img_path, img_name))
        result = table_engine(img)
        save_structure_res(result, save_folder, os.path.basename(img_path).split('.')[0])
        h, w, _ = img.shape
        res = sorted_layout_boxes(result, w)
        convert_info_docx(img, res, save_folder, os.path.basename(img_path).split('.')[0])
    for line in res:
        line.pop('img')
        print(line)
        for pra in line['res']:
            text.append(pra['text'])
            text.append('\n')
    with open('txt/res.txt', 'w', encoding='utf-8') as f:
        f.write('\n'.join(text))
